package com.cloudea.learn.layer;

import com.cloudea.learn.Tensor;

public class ElementPow2 extends Layer {
    @Override
    protected Tensor forward(Context context, Tensor... inputs) {
        Tensor a = inputs[0];
        Tensor c = new Tensor(a.getShape());
        for (int i = 0; i < c.getSize(); i++){
            c.setByLocation(i,  Math.pow(a.getByLocation(i), 2));
        }
        return c;
    }

    @Override
    public Tensor[] backward(Context context, Tensor grid, Tensor... inputs) {
        Tensor a = inputs[0];
        int n = a.getSize();
        Tensor gridA = new Tensor(n);
        for (int i = 0; i < n; i++){
            double deltaA = 2 * a.getByLocation(i);
            gridA.setByLocation(i, grid.getByLocation(i) * deltaA);
        }
        return new Tensor[]{gridA};
    }
}
